# -*- coding: utf-8 -*-
import os
from numpy import *
import numpy as np
import cv2
import glob
import re
import math

# 取消科学计数法
np.set_printoptions(suppress=True)


# 标定
def calib(inter_corner_shape, size_per_grid, img_dir, img_type):
    # criteria: only for subpix calibration, which is not used here.
    # 亚像素比较
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    w, h = inter_corner_shape  # 宽 高 交叉点
    # cp_int: corner point in int form, save the coordinate of corner points in world sapce in 'int' form
    # like (0,0,0), (1,0,0), (2,0,0) ....,(10,7,0).
    # 生成世界坐标 以像素左上角为原点
    cp_int = np.zeros((w * h, 3), np.float32)
    cp_int[:, :2] = np.mgrid[0:w, 0:h].T.reshape(-1, 2)
    # cp_world: corner point in world space, save the coordinate of corner points in world space.
    cp_world = cp_int * size_per_grid

    obj_points = []  # the points in world space # 世界坐标
    # the points in image space (relevant to obj_points) # 图像坐标
    img_points = []
    images = glob.glob(img_dir + os.sep + '**.' + img_type)
    for fname in images:
        # print(re.sub("\D", "", fname))
        img = cv2.imread(fname)
        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # find the corners, cp_img: corner points in pixel space.
        # 寻找内角点
        ret, cp_img = cv2.findChessboardCorners(gray_img, (w, h), None)
        # if ret is True, save.
        if ret == True:
            cv2.cornerSubPix(gray_img, cp_img, (11, 11),
                             (-1, -1), criteria)  # 亚像素级角点检测
            obj_points.append(cp_world)
            img_points.append(cp_img)
            # view the corners
            cv2.drawChessboardCorners(img, (w, h), cp_img, ret)  # 画出内角点
            cv2.imwrite('data/save/p' + re.sub("\D",
                        "", fname) + '.jpg', img)  # 保存
            # cv2.waitKey(1000)
    # cv2.destroyAllWindows()
    # calibrate the camera
    # 标定图片 参数(世界坐标点，图像点坐标，图像大小)
    # 返回值 （ret,内参矩阵,畸变系数,旋转向量,平移向量）
    ret, mat_inter, coff_dis, v_rot, v_trans = cv2.calibrateCamera(obj_points, img_points, gray_img.shape[::-1], None,
                                                                   None)
    print(("ret:"), ret)
    print(("（内参矩阵）internal matrix:\n"), mat_inter)
    # in the form of (k_1,k_2,p_1,p_2,k_3)
    print(("（畸变系数）distortion cofficients:\n"), coff_dis)
    print(("（旋转向量）rotation vectors:\n"), v_rot)
    print(("（平移向量）translation vectors:\n"), v_trans)
    # calculate the error of reproject
    total_error = 0
    for i in range(len(obj_points)):
        # projectPoints()函数用于计算世界坐标系中的三维点投影到像素坐标系中的二维坐标。
        img_points_repro, _ = cv2.projectPoints(
            obj_points[i], v_rot[i], v_trans[i], mat_inter, coff_dis)
        error = cv2.norm(img_points[i], img_points_repro,
                         cv2.NORM_L2) / len(img_points_repro)
        total_error += error
    print(("（ 平均重投影误差）Average Error of Reproject: "),
          total_error / len(obj_points))

    world_points = calcParams(obj_points[0], img_points[0], mat_inter, coff_dis)
    print(world_points)

    return mat_inter, coff_dis


# 矫正图片
def dedistortion(inter_corner_shape, img_dir, img_type, save_dir, mat_inter, coff_dis):
    w, h = inter_corner_shape
    images = glob.glob(img_dir + os.sep + '**.' + img_type)
    for fname in images:
        img_name = fname.split(os.sep)[-1]
        img = cv2.imread(fname)
        newcameramtx, roi = cv2.getOptimalNewCameraMatrix(
            mat_inter, coff_dis, (w, h), 0, (w, h))  # 自由比例参数
        dst = cv2.undistort(img, mat_inter, coff_dis, None, newcameramtx)  # 矫正
        # clip the image
        # x,y,w,h = roi
        # dst = dst[y:y+h, x:x+w]
        cv2.imwrite(save_dir + os.sep + img_name, dst)
    print('Dedistorted images have been saved to: %s successfully.' % save_dir)


def calcParams(points_3D, points_2D, camera_matrix, dist_coeffs):
    # 求相机旋转矩阵 平移矩阵 也就是相机位姿
    '''
    flags:SOLVEPNP_ITERATIVE(迭代法、点在同一平面上)、SOLVEPNP_P3P（3点法）、
    SOLVEPNP_EPNP（引入齐次重心坐标，点可以不在同一平面上）、
    SOLVEPNP_DLS（直接最小二乘，四点不在同一平面）、
    SOLVEPNP_UPNP(尝试估计相机的内在参数)
    ————————————————
    '''
    success, rvec, tvec = cv2.solvePnP(
        points_3D, points_2D, camera_matrix, dist_coeffs, flags=0)
    # 画人脸朝向 
    # 参见 https://www.pythonpool.com/opencv-solvepnp/
    #nose_end_point2D, jacobian = cv2.projectPoints(np.array([(0.0, 0.0, 1000.0)]), rotation_vector, translation_vector, camera_matrix, dist_coeffs)
    # for p in points_2D:
    #     cv2.circle(img, (int(p[0]), int(p[1])), 3, (0,0,255), -1)
    #     point1 = ( int(points_2D[0][0]), int(points_2D[0][1]))
    #     point2 = ( int(nose_end_point2D[0][0][0]), int(nose_end_point2D[0][0][1]))
    #     cv2.line(img, point1, point2, (255,255,255), 2)
    #     # Display image
    #     cv2.imshow(img)
    rotM = cv2.Rodrigues(rvec)[0]
    camera_postion = -np.matrix(rotM).T * np.matrix(tvec)
    # print(camera_postion.T)

    # 验证根据博客http://www.cnblogs.com/singlex/p/pose_estimation_1.html提供方法求解相机位姿
    # 计算相机坐标系的三轴旋转欧拉角，旋转后可以转出世界坐标系。旋转顺序z,y,x
    thetaZ = math.atan2(rotM[1, 0], rotM[0, 0])*180.0/math.pi
    thetaY = math.atan2(-1.0*rotM[2, 0], math.sqrt(rotM[2, 1]
                        ** 2 + rotM[2, 2]**2))*180.0/math.pi
    thetaX = math.atan2(rotM[2, 1], rotM[2, 2])*180.0/math.pi
    # 相机坐标系下值
    x = tvec[0]
    y = tvec[1]
    z = tvec[2]
    (x, y) = RotateByZ(x, y, -1.0*thetaZ)
    (x, z) = RotateByY(x, z, -1.0*thetaY)
    (y, z) = RotateByX(y, z, -1.0*thetaX)
    Cx = x*-1
    Cy = y*-1
    Cz = z*-1

    # 输出相机位置
    print(("输出相机位置 \n"),Cx, Cy, Cz)
    # 输出相机旋转角
    print(("输出相机旋转角 \n"),thetaX, thetaY, thetaZ)
    # 对第五个点进行验证
    Out_matrix = np.concatenate((rotM, tvec), axis=1)
    pixel = np.dot(camera_matrix, Out_matrix)
    pixel1 = np.dot(pixel, np.array([4120, 270, 0, 1], dtype=np.double))
    pixel2 = pixel1/pixel1[2]
    # print(pixel1[2])
    print(pixel2)

    # 计算参数S
    temmat = mat(rotM).I * mat(camera_matrix).I * mat([754, 217, 1]).T
    temmat2 = mat(rotM).I * mat(tvec)
    # print(temmat, temmat2)
    s = temmat2[2]
    s = s/temmat[2]
    # print(s)

    # 计算世界坐标
    wcpoint = mat(rotM).I * (s[0, 0] * mat(camera_matrix).I * mat([754, 217, 1]).T - mat(tvec))
    print(wcpoint)


# 进行三次旋转
def RotateByZ(Cx, Cy, thetaZ):
    rz = thetaZ*math.pi/180.0
    outX = math.cos(rz)*Cx - math.sin(rz)*Cy
    outY = math.sin(rz)*Cx + math.cos(rz)*Cy
    return outX, outY


def RotateByY(Cx, Cz, thetaY):
    ry = thetaY*math.pi/180.0
    outZ = math.cos(ry)*Cz - math.sin(ry)*Cx
    outX = math.sin(ry)*Cz + math.cos(ry)*Cx
    return outX, outZ


def RotateByX(Cy, Cz, thetaX):
    rx = thetaX*math.pi/180.0
    outY = math.cos(rx)*Cy - math.sin(rx)*Cz
    outZ = math.sin(rx)*Cy + math.cos(rx)*Cz
    return outY, outZ




if __name__ == '__main__':
    inter_corner_shape = (8, 6)
    size_per_grid = 0.02  # 方格尺寸 长度单位 米 此处表示20mm

    img_dir = "data/images/2"
    img_type = "jpg"
    # calibrate the camera
    mat_inter, coff_dis = calib(
        inter_corner_shape, size_per_grid, img_dir, img_type)
    # dedistort and save the dedistortion result.
    save_dir = "data/save"
    if (not os.path.exists(save_dir)):
        os.makedirs(save_dir)
    dedistortion(inter_corner_shape, img_dir, img_type,
                 save_dir, mat_inter, coff_dis)
